package com.shen.langchain4j.controller;

import cn.hutool.core.collection.CollUtil;
import dev.langchain4j.data.embedding.Embedding;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.model.output.Response;
import dev.langchain4j.store.embedding.EmbeddingMatch;
import dev.langchain4j.store.embedding.EmbeddingSearchRequest;
import dev.langchain4j.store.embedding.EmbeddingSearchResult;
import dev.langchain4j.store.embedding.EmbeddingStore;
import dev.langchain4j.store.embedding.filter.comparison.IsEqualTo;
import io.qdrant.client.QdrantClient;
import io.qdrant.client.grpc.Collections;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RestController;

import java.util.List;

@Slf4j
@RestController
@RequestMapping("/embedding")
public class ChatEmbeddingController {
    //在Java 15中，文本块现在是官方的，可以使用。
    public static final String POETRY = """
            定风波·莫听穿林打叶声
            （宋）苏轼
            三月七日，沙湖道中遇雨。雨具先去，同行皆狼狈，余独不觉。已而遂晴，故作此词。
            莫听穿林打叶声，何妨吟啸且徐行。
            竹杖芒鞋轻胜马，谁怕？一蓑烟雨任平生。
            料峭春风吹酒醒，微冷，山头斜照却相迎。
            回首向来萧瑟处，归去，也无风雨也无晴。
            """;
    @Autowired
    private EmbeddingModel embeddingModel;
    @Autowired
    private QdrantClient qdrantClient;
    @Autowired
    private EmbeddingStore<TextSegment> embeddingStore;


    /**
     * 通过文本向量化模型进行文本向量化
     *
     * @return 向量化文本
     */
    @GetMapping("/textEmbedding")
    public String textEmbedding() {
        Response<Embedding> embeddingResponse = embeddingModel.embed(POETRY);
        log.info("embeddingResponse: {}", embeddingResponse);
        return embeddingResponse.content().toString();
    }

    /**
     * 新建向量数据库实例和创建索引
     */
    @GetMapping("/createCollection")
    public void createCollection() {
        var vectorParmas = Collections.VectorParams
                .newBuilder()
                .setDistance(Collections.Distance.Cosine)
                .setSize(1024)
                .build();
        qdrantClient.createCollectionAsync("test-qdrant", vectorParmas);
    }

    /**
     * 嵌入存储
     *
     * @return 存储结果
     */
    @GetMapping("/add")
    public String add() {
        TextSegment textSegment = TextSegment.from(POETRY);
        textSegment.metadata().put("author", "苏轼");
        Embedding embedding = embeddingModel.embed(textSegment).content();
        String result = embeddingStore.add(embedding, textSegment);
        log.info("embedding add result: {}", result);
        return result;
    }

    /**
     * 查找最相似的Embedding
     *
     * @return 最相似的Embedding
     */
    @GetMapping("queryOne")
    public String queryOne() {
        Embedding queryEmbedding = embeddingModel.embed("莫听穿林打叶声").content();
        EmbeddingSearchRequest embeddingSearchRequest = EmbeddingSearchRequest.builder()
                .queryEmbedding(queryEmbedding)
                .maxResults(1)
                .build();
        EmbeddingSearchResult<TextSegment> searchResult = embeddingStore.search(embeddingSearchRequest);
        List<EmbeddingMatch<TextSegment>> matchList = searchResult.matches();
        if (CollUtil.isEmpty(matchList)) {
            return "未找到相似的Embedding";
        }
        String resultText = searchResult.matches().getFirst().embedded().text();
        log.info("searchResult: {}", resultText);
        return resultText;
    }

    /**
     * 查找最相似的Embedding（带过滤器Filter）
     *
     * @return 最相似的Embedding
     */
    @GetMapping("queryTwo")
    public String queryTwo() {
        Embedding queryEmbedding = embeddingModel.embed("莫听穿林打叶声").content();
        EmbeddingSearchRequest embeddingSearchRequest = EmbeddingSearchRequest.builder()
                .queryEmbedding(queryEmbedding)
                .filter(new IsEqualTo("author", "王维"))
                .maxResults(1)
                .build();
        EmbeddingSearchResult<TextSegment> searchResult = embeddingStore.search(embeddingSearchRequest);
        List<EmbeddingMatch<TextSegment>> matchList = searchResult.matches();
        if (CollUtil.isEmpty(matchList)) {
            return "未找到相似的Embedding";
        }
        String resultText = searchResult.matches().getFirst().embedded().text();
        log.info("searchResult: {}", resultText);
        return resultText;
    }
}
